//******************************************************************************
// cmix_h5.cc - A BASF module for reading MDST files and extracting 
// quantities of interest for events consistent with a mu+mu-pi0pi0 final state.
//
// Date: June 2012
// Author: craigb (craig dot bookwalter at pnnl dot gov)
//******************************************************************************

// Standard headers
#include <algorithm>
#include <exception>
#include <iostream>
#include <fstream>
#include <set>
#include <stdexcept>
#include <string>
#include <sstream>
#include <utility>
#include <vector>

// Belle-specific headers
#include "particle/Particle.h"
#include "particle/PID.h"
#include "particle/utility.h"
#include "event/BelleEvent.h"
#include "basf/module.h"
#include "basf/module_descr.h"
#include "benergy/BeamEnergy.h"
#include "kid/atc_pid.h"
#include "ip/IpProfile.h"
#include "helix/Helix.h"
#include "kfitter/kvertexfitter.h"

#include <mdst/mdst.h>

#include "panther/panther.h"
#include MDST_H
#include MDST_OBS_H
#include EVTCLS_H
#include BELLETDF_H
#include RUN_INFO_H

// Headers written by me
#include "rows/cmix_countrow.hpp"
#include "rows/cmix_chargedrow.hpp"

#include <pytpp.hpp>
#include <csbbelle.hpp>

using std::string;
using std::vector;
using std::cout;
using std::cerr;
using std::endl;

static const double RADS_TO_DEG = 360./(2*3.14159265359);
static const double D0_MASS= 1.86486;
static const double D0_SIGMA=0.01;


#ifdef BELLE_NAMESPACE
namespace Belle {
#endif //BELLE_NAMESPACE

  // All BASF modules inherit from the "Module" class, defined in `basf/module.h'
  // They are written as a single class. The BASF executable loads this class
  // (compiled into a shared object) dynamically. The module constructor, `init' 
  // and `hist_def' methods are all called before reading events from the event 
  // source. The first event is read and BASF calls `begin_run()'. Every time 
  // the experimental run number (contained in the datastream) changes, `end_run()' 
  // will be called and then `begin_run()' again.
  //
  // The `event()' function does the event-by-event processing; in the case of 
  // this module, saving some quantities of interest to an HDF5 dataset for 
  // later study. 
  //
  // When the event source is exhausted, the term() function is called, followed 
  // by the class destructor.
  //----------------------------------------------------------------------------
  class cmix_h5: public Module {
  public:
    cmix_h5();
    virtual ~cmix_h5();
    void init(int *);
    void term();
    void disp_stat(const char*);
    void hist_def();
    void event(BelleEvent*, int*);
    void begin_run(BelleEvent*, int*);
    void end_run(BelleEvent*, int*);

    char outpath_stub[4096];
    char inpath[4096]; 

    static HepPoint3D IP;
    static HepSymMatrix IPerr;

  private:
    bool _write_h5;
    bool _write_hbook;

    int _nevents_tot;
    int _combos_found;

    pyt::table<cmix_chargedrow_t>* _nt; 
    cmix_chargedrow_t* _row;

    pyt::table<cmix_countrow_t>* _countnt; 
    cmix_countrow_t* _countrow;
  };

  HepPoint3D cmix_h5::IP = HepPoint3D(0,0,0);
  HepSymMatrix cmix_h5::IPerr = HepSymMatrix(3,0);


  //----------------------------------------------------------------------------
  // This function registers the module with the BASF executeable.
  // Any parameters defined in the control script (`cmix_h5.basf') are
  // declared and matched with public member variables of the module class.
  //----------------------------------------------------------------------------
  extern "C" Module_descr *mdcl_cmix_h5()
  {

    cmix_h5 *tmp = new cmix_h5();
    Module_descr *dscr = new Module_descr ( "cmix_h5", tmp);
    dscr->define_param("outpath_stub", "Path stub to output file", "S", 4096, tmp->outpath_stub);
    dscr->define_param("inpath", "Input file", "S", 4096, tmp->inpath);

    BeamEnergy::define_global(dscr);
    IpProfile::define_global(dscr);

    return dscr;
  }
  
  //----------------------------------------------------------------------------
  // Default constructor
  //----------------------------------------------------------------------------
  cmix_h5::cmix_h5() : 
  _combos_found(0), 
    _row(new cmix_chargedrow_t),
    _countrow(new cmix_countrow_t),
    _nt(NULL),
    _countnt(NULL)
  {
    // Initialize member vars
  }

  //----------------------------------------------------------------------------
  // Default destructor
  //----------------------------------------------------------------------------  
  cmix_h5::~cmix_h5() {
    delete _nt;
    delete _row;
    delete _countnt;
    delete _countrow;
  }

  //----------------------------------------------------------------------------
  // In init(), I define one or more output files based on the values of the 
  // parameters passed in from the control scripts.
  //----------------------------------------------------------------------------    
  void cmix_h5::init(int* status) {
    string path_stub(outpath_stub);


    string outpath = path_stub + ".ntcuts.h5";
      
    _nt = pyt::open_table<cmix_chargedrow_t>(outpath, "nt", "w", 
					     "Ntuple for charm-mixing analysis",
					     0);
    
    _nt->table_setattr("mdst_path", inpath);
    _nt->table_setattr("nttype", "allcombos");
    _nt->table_setattr("module", "cmix");
    
    outpath = path_stub + ".count.h5";

    _countnt = pyt::open_table<cmix_countrow_t>(outpath, "nt", "w", 
						"Event-counting ntuple for charm-mixing analysis",
						0);
    
    _countnt->table_setattr("mdst_path", inpath);
    _countnt->table_setattr("module", "cmix");
    
    cout << "CRAIGB : Analyzing file " << inpath << endl;
  }
  
  //----------------------------------------------------------------------------
  // hist_def() I leave empty
  //----------------------------------------------------------------------------  
  void cmix_h5::hist_def() {
    // Define ntuples here
  }


  //----------------------------------------------------------------------------
  // begin_run() just contains some bookkeeping to make sure the numbers I save
  // are well-calibrated.
  //----------------------------------------------------------------------------    
  void cmix_h5::begin_run(BelleEvent* be, int* status) {
    BeamEnergy::begin_run();
    IpProfile::begin_run();
  }
 

  //----------------------------------------------------------------------------
  // h5event() is called if the `h5' output format is requested from the
  // control script; again, it processes a single event and saves quantities of 
  // interest to the output file.
  //----------------------------------------------------------------------------  
  void cmix_h5::event(BelleEvent* evptr, int* status) {
    Belle_event_Manager& m_event = Belle_event_Manager::get_manager();

    _nevents_tot++;

    _countrow->expno   = m_event[0].ExpNo();
    _countrow->runno   = m_event[0].RunNo();
    _countrow->eventno = m_event[0].EvtNo();

    _countrow->kept   = 1;
    _countrow->npos  = -1;
    _countrow->nneg = -1;
    _countrow->ngamma = -1;
    _countrow->npi0   = -1;
    
    Mdst_charged_Manager& m_charged = Mdst_charged_Manager::get_manager();
    _countrow->ncharged_raw = m_charged.count();
    
    if (IpProfile::usable()) {
      IP    = IpProfile::position(1);
      IPerr = IpProfile::position_err_b_life_smeared(1);    
    }
    else {
      IP    = HepPoint3D(0,0,0);
      IPerr = HepSymMatrix(3,0);
    }
            
    // Build charged track combos
    vector<Mdst_charged> vpos;
    vector<Mdst_charged> vneg;

    for (Mdst_charged_Manager::iterator it = m_charged.begin();
	 it != m_charged.end(); it++) 
      {
	if (!good_charged(*it))
	  continue;
	
	if (it->charge() > 0)
	  vpos.push_back(*it);
	else
	  vneg.push_back(*it);
      }
    _countrow->npos = vpos.size();
    _countrow->nneg = vneg.size();
                
    HepLorentzVector initial = initial_state(); 
    HepLorentzVector initial_cm = initial_state();

    Hep3Vector boost_cm      = initial.boostVector();   
    initial_cm.boost(-boost_cm);

    _countnt->append(*_countrow);

    atc_pid pid_pik(3,1,5,2,3);
    atc_pid pid_kpi(3,1,5,3,2);

    for (vector<Mdst_charged>::iterator ipos = vpos.begin(); ipos != vpos.end(); ipos++) {
      for (vector<Mdst_charged>::iterator ineg = vneg.begin(); ineg != vneg.end(); ineg++) {	
	_row->expno   = m_event[0].ExpNo();
	_row->runno   = m_event[0].RunNo();
	_row->eventno = m_event[0].EvtNo();
	_row->beam_e  = BeamEnergy::E_beam_corr();

	_row->idhash = two_body_idhash(ipos->get_ID(), ineg->get_ID());

	_row->pos_kpi_llh = pid_kpi.prob(*ipos);
	_row->pos_pik_llh = pid_pik.prob(*ipos);

	_row->pos_px = ipos->px();
	_row->pos_py = ipos->py();
	_row->pos_pz = ipos->pz();

	_row->neg_kpi_llh = pid_kpi.prob((*ineg));
	_row->neg_pik_llh = pid_pik.prob((*ineg));

	_row->neg_px = ineg->px();
	_row->neg_py = ineg->py();
	_row->neg_pz = ineg->pz();
	
	HepLorentzVector p4pos_pi = charged_track_p4(IP, *ipos, 2);
	HepLorentzVector p4pos_k  = charged_track_p4(IP, *ipos, 3);

	HepLorentzVector p4neg_pi = charged_track_p4(IP, *ineg, 2);
	HepLorentzVector p4neg_k  = charged_track_p4(IP, *ineg, 3);
		
	_row->m_kk   = (p4pos_k + p4neg_k).m();
	_row->m_kpi  = (p4pos_k + p4neg_pi).m();
	_row->m_pik  = (p4pos_pi + p4neg_k).m();
	_row->m_pipi = (p4pos_pi + p4neg_pi).m();
	
	// if (fabs(_row->m_kk - D0_MASS) < 2*D0_SIGMA)
	//   continue;

	// if (fabs(_row->m_pipi - D0_MASS) < 2*D0_SIGMA)
	//   continue;

	_nt->append(*_row);
	_combos_found++;
      }
    }
  }

  

  
  //----------------------------------------------------------------------------
  // end_run() handles any bookkeeping necessary at the end of an experimental
  // run--in this case it's empty.
  //----------------------------------------------------------------------------    
  void cmix_h5::end_run(BelleEvent* be, int* status) { 
    // Yessir
  }

  //----------------------------------------------------------------------------
  // term() is used to write out any output waiting to be flushed
  //----------------------------------------------------------------------------    
  void cmix_h5::term() {
    cout << "CRAIGB : " << _nevents_tot << " events processed, " <<  _combos_found << " D0/D0bar candidates found" << endl;    
    
    // End of analysis, close ntuples
  }


  //----------------------------------------------------------------------------
  // disp_stat() would display some statistics about the events I saved or threw
  // away if I was motivated enough to keep track of such things.
  //----------------------------------------------------------------------------    
  void cmix_h5::disp_stat(const char* str) {
    // Who knows
  }
                   
  //----------------------------------------------------------------------------
  // calc_R2() is a function that calculates a physics quantity (normalized Fox-
  // Wolfram second moment, if you're curious).
  //----------------------------------------------------------------------------    
    

#ifdef BELLE_NAMESPACE
}
#endif
